// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/dwrf/reader/selective_repeated_column_reader.h>
#include <pollux/dwio/dwrf/reader/selective_dwrf_reader.h>

namespace kumo::pollux::dwrf {
namespace {
std::unique_ptr<dwio::common::IntDecoder</*isSigned*/ false>> makeLengthDecoder(
    const dwio::common::TypeWithId& fileType,
    DwrfParams& params,
    memory::MemoryPool& pool) {
  EncodingKey encodingKey{fileType.id(), params.flatMapContext().sequence};
  auto& stripe = params.stripeStreams();
  const auto rleVersion = convertRleVersion(stripe, encodingKey);
  const auto lenId = StripeStreamsUtil::getStreamForKind(
      stripe,
      encodingKey,
      proto::Stream_Kind_LENGTH,
      proto::orc::Stream_Kind_LENGTH);
  const bool lenVints = stripe.getUseVInts(lenId);
  return createRleDecoder</*isSigned=*/false>(
      stripe.getStream(lenId, params.streamLabels().label(), true),
      rleVersion,
      pool,
      lenVints,
      dwio::common::INT_BYTE_SIZE);
}
} // namespace

FlatMapContext flatMapContextFromEncodingKey(const EncodingKey& encodingKey) {
  return FlatMapContext{
      .sequence = encodingKey.sequence(),
      .inMapDecoder = nullptr,
      .keySelectionCallback = nullptr};
}

SelectiveListColumnReader::SelectiveListColumnReader(
    const TypePtr& requestedType,
    const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
    DwrfParams& params,
    common::ScanSpec& scanSpec)
    : dwio::common::SelectiveListColumnReader(
          requestedType,
          fileType,
          params,
          scanSpec),
      length_(makeLengthDecoder(*fileType_, params, *memoryPool_)) {
  POLLUX_CHECK_EQ(fileType_->id(), fileType->id(), "working on the same node");
  EncodingKey encodingKey{fileType_->id(), params.flatMapContext().sequence};
  auto& stripe = params.stripeStreams();
  // count the number of selected sub-columns
  auto& childType = requestedType_->childAt(0);
  if (scanSpec_->children().empty()) {
    scanSpec.getOrCreateChild(common::ScanSpec::kArrayElementsFieldName);
  }
  scanSpec_->children()[0]->setProjectOut(true);

  auto childParams = DwrfParams(
      stripe,
      params.streamLabels(),
      params.runtimeStatistics(),
      flatMapContextFromEncodingKey(encodingKey));
  child_ = SelectiveDwrfReader::build(
      childType, fileType_->childAt(0), childParams, *scanSpec_->children()[0]);
  children_ = {child_.get()};
}

SelectiveMapColumnReader::SelectiveMapColumnReader(
    const TypePtr& requestedType,
    const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
    DwrfParams& params,
    common::ScanSpec& scanSpec)
    : dwio::common::SelectiveMapColumnReader(
          requestedType,
          fileType,
          params,
          scanSpec),
      length_(makeLengthDecoder(*fileType_, params, *memoryPool_)) {
  POLLUX_CHECK_EQ(fileType_->id(), fileType->id(), "working on the same node");
  const EncodingKey encodingKey{
      fileType_->id(), params.flatMapContext().sequence};
  auto& stripe = params.stripeStreams();
  if (scanSpec_->children().empty()) {
    scanSpec_->getOrCreateChild(common::ScanSpec::kMapKeysFieldName);
    scanSpec_->getOrCreateChild(common::ScanSpec::kMapValuesFieldName);
  }
  scanSpec_->children()[0]->setProjectOut(true);
  scanSpec_->children()[1]->setProjectOut(true);

  auto& keyType = requestedType_->childAt(0);
  auto keyParams = DwrfParams(
      stripe,
      params.streamLabels(),
      params.runtimeStatistics(),
      flatMapContextFromEncodingKey(encodingKey));
  keyReader_ = SelectiveDwrfReader::build(
      keyType,
      fileType_->childAt(0),
      keyParams,
      *scanSpec_->children()[0].get());

  auto& valueType = requestedType_->childAt(1);
  auto elementParams = DwrfParams(
      stripe,
      params.streamLabels(),
      params.runtimeStatistics(),
      flatMapContextFromEncodingKey(encodingKey));
  elementReader_ = SelectiveDwrfReader::build(
      valueType,
      fileType_->childAt(1),
      elementParams,
      *scanSpec_->children()[1]);
  children_ = {keyReader_.get(), elementReader_.get()};
}

} // namespace kumo::pollux::dwrf
